# -*- coding: utf-8 -*-
# @Time    : 2024/7/3 10:12
# @Author  : yujiahao
# @File    : 03_threadpool_async.py
# @description:将数据库表的数据导出到文件系统（普通写，多线程、多进程、异步io）
import logging
import pymysql
from dbutils.pooled_db import PooledDB
import asyncio
import aiomysql
import aiofiles
from concurrent.futures import ThreadPoolExecutor
from multiprocessing import Process, Queue, current_process
import time
from functools import wraps

"""
解题思路：
1. 配置数据库连接池：
    - 设置数据库连接池的参数，包括主机、用户、密码、数据库、自动提交、最小连接数和最大连接数等。
    - 初始化连接池，以便后续数据库操作可以复用连接，提升效率。

2. 封装同步和异步的写文件函数：
    - 同步写文件函数：打开文件并逐条写入数据。
    - 异步写文件函数：使用 aiofiles 异步写入数据，避免阻塞事件循环。

3. 实现单条写出功能：
    - 从连接池中获取一个连接。
    - 使用游标执行查询。
    - 分批次获取查询结果并写入文件。
    - 释放连接回连接池。

4. 实现批量写出功能：
    - 从连接池中获取一个连接。
    - 使用游标执行查询。
    - 分批次获取查询结果并写入文件。
    - 释放连接回连接池。

5. 实现使用线程池异步写出功能：
    - 定义一个函数 fetch_and_write，从连接池中获取连接并执行查询。
    - 使用 ThreadPoolExecutor 创建线程池并提交任务。
    - 分批次获取查询结果并写入文件。
    - 释放连接回连接池。

6. 实现使用 asyncio 异步写出功能：
    - 定义一个异步函数 fetch_and_write_async，使用 aiomysql 创建异步数据库连接池。
    - 从连接池中获取连接并执行查询。
    - 分批次获取查询结果并异步写入文件。
    - 关闭连接池。

7. 实现使用多进程写出功能：
    - 使用 multiprocessing 模块创建多个进程。
    - 每个进程从连接池中获取连接并执行查询。
    - 分批次获取查询结果并写入文件。
"""

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


# 配置数据库连接池
def init_db_pool():
    """
    初始化数据库连接池。

    Returns:
        PooledDB: 数据库连接池对象。
    """
    pool = PooledDB(
        creator=pymysql,  # 创建数据库连接的模块
        maxconnections=20,  # 连接池允许的最大连接数
        mincached=5,  # 初始化时创建的连接数
        maxcached=10,  # 连接池中最多闲置的连接数
        blocking=True,  # 连接池中没有可用连接后是否阻塞等待
        maxusage=None,  # 单个连接的最大重复使用次数
        setsession=['SET AUTOCOMMIT = 1'],  # 会话参数
        ping=1,  # 检查连接是否可用
        host='127.0.0.1',
        port=33061,
        user='***',
        passwd='***',
        db='bigdata_test',
        charset='utf8'
    )
    return pool


# 初始化数据库连接池
db_pool = init_db_pool()


def measure_time(func):
    """
    装饰器函数，用于测量函数执行时间。

    Args:
        func (function): 被装饰的函数。

    Returns:
        function: 包装后的函数。
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        logging.info(f"开始执行函数 {func.__name__}")
        result = func(*args, **kwargs)
        end_time = time.time()
        logging.info(f"函数 {func.__name__} 执行时间: {end_time - start_time:.2f} 秒")
        return result

    return wrapper


def measure_time_async(func):
    """
    异步装饰器函数，用于测量异步函数执行时间。

    Args:
        func (function): 被装饰的异步函数。

    Returns:
        function: 包装后的异步函数。
    """

    @wraps(func)
    async def wrapper(*args, **kwargs):
        start_time = time.time()
        logging.info(f"开始执行异步函数 {func.__name__}")
        result = await func(*args, **kwargs)
        end_time = time.time()
        logging.info(f"异步函数 {func.__name__} 执行时间: {end_time - start_time:.2f} 秒")
        return result

    return wrapper


@measure_time
def write_to_file(filename, rows):
    """
    同步将查询结果写入文件。

    Args:
        filename (str): 要写入的文件名。
        rows (list): 查询结果行列表。
    """
    logging.info(f"开始将数据写入文件 {filename}")
    with open(filename, 'a') as f:
        for row in rows:
            f.write(str(row) + '\n')
    logging.info(f"完成将数据写入文件 {filename}")


@measure_time_async
async def write_to_file_async(filename, rows):
    """
    异步将查询结果写入文件。

    Args:
        filename (str): 要写入的文件名。
        rows (list): 查询结果行列表。
    """
    logging.info(f"开始将数据异步写入文件 {filename}")
    async with aiofiles.open(filename, 'a') as f:
        for row in rows:
            await f.write(str(row) + '\n')
    logging.info(f"完成将数据异步写入文件 {filename}")


def fetch_and_write(conn, cursor, batch_size, filename, sql):
    """
    从数据库中获取数据并写入文件。

    Args:
        conn (pymysql.connections.Connection): 数据库连接对象。
        cursor (pymysql.cursors.Cursor): 数据库游标对象。
        batch_size (int): 每次获取的行数。
        filename (str): 要写入的文件名。
        sql (str): 要执行的 SQL 查询语句。
    """
    logging.info(f"开始从数据库中获取数据并写入文件 {filename}")
    cursor.execute(sql)
    while True:
        rows = cursor.fetchmany(batch_size)
        if not rows:
            break
        write_to_file(filename, rows)
    logging.info(f"完成从数据库中获取数据并写入文件 {filename}")


@measure_time
def write_single(sql, filename, batch_size=10000):
    """
    单条写出：从数据库中逐条读取数据并写入文件。

    Args:
        sql (str): 要执行的 SQL 查询语句。
        filename (str): 要写入的文件名。
        batch_size (int, optional): 每次获取的行数。默认值为10000。
    """
    logging.info("开始单条写出")
    conn = db_pool.connection()  # 获取数据库连接
    try:
        with conn.cursor() as cursor:
            fetch_and_write(conn, cursor, batch_size, filename, sql)
    finally:
        conn.close()  # 关闭数据库连接
    logging.info("完成单条写出")


@measure_time
def write_batch(sql, filename, batch_size=10000):
    """
    批量写出：从数据库中按批次读取数据并写入文件。

    Args:
        sql (str): 要执行的 SQL 查询语句。
        filename (str): 要写入的文件名。
        batch_size (int, optional): 每次获取的行数。默认值为10000。
    """
    logging.info("开始批量写出")
    conn = db_pool.connection()  # 获取数据库连接
    try:
        with conn.cursor() as cursor:
            fetch_and_write(conn, cursor, batch_size, filename, sql)
    finally:
        conn.close()  # 关闭数据库连接
    logging.info("完成批量写出")


@measure_time
def write_with_threadpool(sql, filename, batch_size=10000):
    """
    使用线程池异步写出：利用线程池从数据库中读取数据并写入文件。

    Args:
        sql (str): 要执行的 SQL 查询语句。
        filename (str): 要写入的文件名。
        batch_size (int, optional): 每次获取的行数。默认值为10000。
    """
    logging.info("开始使用线程池异步写出")

    def fetch_and_write_in_thread(sql, filename, batch_size):
        conn = db_pool.connection()  # 获取数据库连接
        try:
            with conn.cursor() as cursor:
                fetch_and_write(conn, cursor, batch_size, filename, sql)
        finally:
            conn.close()  # 关闭数据库连接

    with ThreadPoolExecutor() as executor:
        future = executor.submit(fetch_and_write_in_thread, sql, filename, batch_size)
        # 等待任务完成
        future.result()

    logging.info("完成使用线程池异步写出")


@measure_time_async
async def fetch_and_write_async(pool, batch_size, filename, sql):
    """
    异步从数据库中获取数据并写入文件。

    Args:
        pool (aiomysql.Pool): 异步数据库连接池。
        batch_size (int): 每次获取的行数。
        filename (str): 要写入的文件名。
        sql (str): 要执行的 SQL 查询语句。
    """
    logging.info(f"开始异步从数据库中获取数据并写入文件 {filename}")
    async with pool.acquire() as conn:
        async with conn.cursor() as cursor:
            await cursor.execute(sql)
            while True:
                rows = await cursor.fetchmany(batch_size)
                if not rows:
                    break
                await write_to_file_async(filename, rows)
    logging.info(f"完成异步从数据库中获取数据并写入文件 {filename}")


@measure_time_async
async def write_with_asyncio(sql, filename, batch_size=10000):
    """
    使用 asyncio 异步写出：利用 asyncio 从数据库中异步读取数据并写入文件。

    Args:
        sql (str): 要执行的 SQL 查询语句。
        filename (str): 要写入的文件名。
        batch_size (int, optional): 每次获取的行数。默认值为10000。
    """
    logging.info("开始使用 asyncio 异步写出")
    pool = await aiomysql.create_pool(
        host='127.0.0.1',
        port=33061,
        user='***',
        passwd='***',
        db='bigdata_test',
        minsize=1,
        maxsize=10
    )
    await fetch_and_write_async(pool, batch_size, filename, sql)
    pool.close()
    await pool.wait_closed()
    logging.info("完成使用 asyncio 异步写出")


def worker(queue, batch_size, sql):
    """
    多进程工作函数：从数据库中读取数据并将其放入队列。

    Args:
        queue (Queue): 用于在进程间传递数据的队列。
        batch_size (int): 每次获取的行数。
        sql (str): 要执行的 SQL 查询语句。
    """
    logging.info(f"多进程工作函数开始，进程ID: {current_process().pid}")
    conn = pymysql.connect(
        host='127.0.0.1',
        port=33061,
        user='***',
        passwd='***',
        db='bigdata_test',
        charset='utf8'
    )
    try:
        with conn.cursor() as cursor:
            cursor.execute(sql)
            while True:
                rows = cursor.fetchmany(batch_size)
                if not rows:
                    break
                queue.put(rows)
    finally:
        conn.close()  # 关闭数据库连接
    logging.info(f"多进程工作函数结束，进程ID: {current_process().pid}")


@measure_time
def write_with_multiprocessing(sql, filename, batch_size=10000, num_processes=2):
    """
    使用多进程写出：利用多进程从数据库中读取数据并写入文件。

    Args:
        sql (str): 要执行的 SQL 查询语句。
        filename (str): 要写入的文件名。
        batch_size (int, optional): 每次获取的行数。默认值为10000。
        num_processes (int, optional): 启动的进程数。默认值为2。
    """
    logging.info("开始使用多进程写出")
    processes = []
    queue = Queue()

    # 启动多个进程，每个进程从数据库中读取数据并放入队列
    for _ in range(num_processes):
        p = Process(target=worker, args=(queue, batch_size, sql))
        p.start()
        processes.append(p)

    # 主进程从队列中获取数据并写入文件
    with open(filename, 'a') as f:
        while any(p.is_alive() for p in processes) or not queue.empty():
            while not queue.empty():
                rows = queue.get()
                for row in rows:
                    f.write(str(row) + '\n')

    # 等待所有进程结束
    for p in processes:
        p.join()

    logging.info("完成使用多进程写出")


def main():
    sql = "SELECT * FROM bigdata_test.user_app_test limit 10000"

    # 定义文件名变量
    output_single = "output_single.txt"
    output_batch = "output_batch.txt"
    output_threadpool = "output_threadpool.txt"
    output_asyncio = "output_asyncio.txt"
    output_multiprocessing = "output_multiprocessing.txt"

    # 单条写出示例
    write_single(sql, output_single)

    # 批量写出示例
    write_batch(sql, output_batch)

    # 使用线程池异步写出示例
    write_with_threadpool(sql, output_threadpool)

    # 使用 asyncio 异步写出示例
    asyncio.run(write_with_asyncio(sql, output_asyncio))

    # 使用多进程写出示例
    write_with_multiprocessing(sql, output_multiprocessing)


if __name__ == "__main__":
    main()
